package id.mankel.commons.util;

import org.springframework.core.task.TaskExecutor;

import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;
import java.util.stream.Stream;

public class MkStreams {

    public static <T, R> Stream<R> parallelMap(TaskExecutor taskExecutor, Collection<T> c, Function<T, R> map) {
        return parallelMap(taskExecutor, c.stream(), map);
    }

    public static <T, R> Stream<R> parallelMap(TaskExecutor taskExecutor, Stream<T> s, Function<T, R> map) {
        List<CompletableFuture<R>> l = s.map(t -> CompletableFuture.supplyAsync(() -> map.apply(t), taskExecutor)).toList();
        CompletableFuture.allOf(l.toArray(new CompletableFuture[0])).join();
        return l.stream().map(CompletableFuture::join);
    }

    public static <T> void parallelForEach(TaskExecutor taskExecutor, Collection<T> c, Runnable runnable) {
        CompletableFuture.allOf(c.stream().map(t -> CompletableFuture.runAsync(runnable, taskExecutor)).toArray(CompletableFuture[]::new)).join();
    }

}
